{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.2. 张量程序抽象\n",
    "2.2.1. 张量程序抽象中的其它结构\n",
    "张量程序可以将这些额外的信息合并为程序的一部分，以使程序变换更加便利。\n",
    "\n",
    "举个例子，上面图中的程序包含额外的 T.axis.spatial 标注，表明 vi 这个特定的变量被映射到循环变量 i，并且所有的迭代都是独立的。这个信息对于执行这个程序而言并非必要，但会使得我们在变换这个程序时更加方便。在这个例子中，我们知道我们可以安全地并行或者重新排序所有与 vi 有关的循环，只要实际执行中 vi 的值按照从 0 到 128 的顺序变化。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mlc.ai/zh/chapter_tensor_program/case_study.html\n",
    "2.4. TensorIR: 张量程序抽象案例研究¶\n",
    "2.4.3. 学习一类张量程序抽象 – TensorIR¶\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = \"float32\"\n",
    "a_np = np.random.rand(128, 128).astype(dtype)\n",
    "b_np = np.random.rand(128, 128).astype(dtype)\n",
    "# a @ b is equivalent to np.matmul(a, b)\n",
    "c_mm_relu = np.maximum(a_np @ b_np, 0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在看过低级 NumPy 示例后，现在我们准备介绍 TensorIR。下面的代码块展示了 mm_relu 的 TensorIR 实现。这里的代码是用一种名为 TVMScript 的语言实现的，它是一种嵌入在 Python AST 中的特定领域方言。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def mm_relu(A: T.Buffer((128, 128), \"float32\"),\n",
    "                B: T.Buffer((128, 128), \"float32\"),\n",
    "                C: T.Buffer((128, 128), \"float32\")):\n",
    "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
    "        Y = T.alloc_buffer((128, 128), dtype=\"float32\")\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            with T.block(\"Y\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                vk = T.axis.reduce(128, k)\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
    "        for i, j in T.grid(128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.3.5. 为什么块需要额外附加的信息\n",
    "一个重要的观察结果是，附加信息（块轴范围及其属性）使块轴独立于外部循环嵌套 i, j, k。\n",
    "\n",
    "块轴信息还提供了额外的属性，帮助我们验证用于执行计算的外部循环的正确性。例如，下面的代码块会导致错误，因为循环需要一个大小为 128 的迭代器，但我们只将它绑定到一个大小为 127 的 for 循环。\n",
    "\n",
    "# wrong program due to loop and block iteration mismatch\n",
    "for i in range(127):\n",
    "    with T.block(\"C\"):\n",
    "        vi = T.axis.spatial(128, i)\n",
    "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
    "        error here due to iterator size mismatch\n",
    "        ...\n",
    "Copy to clipboard\n",
    "这些附加信息也有助于我们进行机器学习编译分析。例如，虽然我们总是可以在空间轴上做并行化，在规约轴上进行并行化将需要特定的策略。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.4. 变换\n",
    "在上一节中，我们了解了 TensorIR 及其关键元素。 现在，让我们了解所有机器学习编译工作流的主要成分——元张量函数的变换。\n",
    "\n",
    "在上一节中，我们给出了如何使用低级 NumPy 编写 mm_relu 的示例。 在实践中，可以有多种方法来实现相同的功能，并且每种实现都可能导致不同的性能。\n",
    "\n",
    "我们将讨论性能背后的原因以及如何在以后的内容中利用这些变体。 在本节中，让我们关注使用变换获得不同实现变体的能力。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">,</span> <span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">,</span> <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SSR&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{,} \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SSR}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"tir.noalias\": True, \"global_symbol\": \"mm_relu\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j, k in T.grid(128, 128, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import IPython\n",
    "IPython.display.Code(MyModule.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.3.6. 块轴绑定的语法糖\n",
    "在每个块轴直接映射到外部循环迭代器的情况下，我们可以使用 T.axis.remap 在一行中声明所有块轴。\n",
    "\n",
    "# SSR means the properties of each axes are \"spatial\", \"spatial\", \"reduce\"\n",
    "vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "Copy to clipboard\n",
    "相当于\n",
    "\n",
    "vi = T.axis.spatial(range_of_i, i)\n",
    "vj = T.axis.spatial(range_of_j, j)\n",
    "vk = T.axis.reduce(range_of_k, k)\n",
    "Copy to clipboard\n",
    "所以我们也可以编写如下程序。\n",
    "\n",
    "@tvm.script.ir_module\n",
    "class MyModuleWithAxisRemapSugar:\n",
    "    @T.prim_func\n",
    "    def mm_relu(A: T.Buffer((128, 128), \"float32\"),\n",
    "                B: T.Buffer((128, 128), \"float32\"),\n",
    "                C: T.Buffer((128, 128), \"float32\")):\n",
    "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
    "        Y = T.alloc_buffer((128, 128), dtype=\"float32\")\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            with T.block(\"Y\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
    "        for i, j in T.grid(128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.4. 变换\n",
    "\n",
    "在上一节中，我们了解了 TensorIR 及其关键元素。 现在，让我们了解所有机器学习编译工作流的主要成分——元张量函数的变换。\n",
    "\n",
    "在上一节中，我们给出了如何使用低级 NumPy 编写 mm_relu 的示例。 在实践中，可以有多种方法来实现相同的功能，并且每种实现都可能导致不同的性能。\n",
    "\n",
    "我们将讨论性能背后的原因以及如何在以后的内容中利用这些变体。 在本节中，让我们关注使用变换获得不同实现变体的能力。\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):\n",
    "    Y = np.empty((128, 128), dtype=\"float32\")\n",
    "    for i in range(128):\n",
    "        for j0 in range(32):\n",
    "            for k in range(128):\n",
    "                for j1 in range(4):\n",
    "                    j = j0 * 4 + j1\n",
    "                    if k == 0:\n",
    "                        Y[i, j] = 0\n",
    "                    Y[i, j] = Y[i, j] + A[i, k] * B[k, j]\n",
    "    for i in range(128):\n",
    "        for j in range(128):\n",
    "            C[i, j] = max(Y[i, j], 0)\n",
    "\n",
    "c_np = np.empty((128, 128), dtype=dtype)\n",
    "lnumpy_mm_relu_v2(a_np, b_np, c_np)\n",
    "np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">,</span> <span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span><span class=\"p\">,</span> <span class=\"n\">j_1</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0}\\PY{p}{,} \\PY{n}{j\\PYZus{}1}\\PY{p}{,} \\PY{n}{k} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"tir.noalias\": True, \"global_symbol\": \"mm_relu\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0, j_1, k in T.grid(128, 32, 4, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                vi = T.axis.spatial(128, i)\n",
       "                vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                vk = T.axis.reduce(128, k)\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(MyModule)\n",
    "block_Y = sch.get_block(\"Y\", func_name=\"mm_relu\")\n",
    "i, j, k = sch.get_loops(block_Y)\n",
    "j0, j1 = sch.split(j, factors=[None, 4])\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">,</span> <span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0}\\PY{p}{,} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"tir.noalias\": True, \"global_symbol\": \"mm_relu\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0, k, j_1 in T.grid(128, 32, 128, 4):\n",
       "            with T.block(\"Y\"):\n",
       "                vi = T.axis.spatial(128, i)\n",
       "                vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                vk = T.axis.reduce(128, k)\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sch.reorder(j0, k, j1)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">,</span> <span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                        <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">ax0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">ax0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{for} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                    \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                        \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "            \\PY{k}{for} \\PY{n}{ax0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{ax0}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"tir.noalias\": True, \"global_symbol\": \"mm_relu\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0 in T.grid(128, 32):\n",
       "            for k, j_1 in T.grid(128, 4):\n",
       "                with T.block(\"Y\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                    vk = T.axis.reduce(128, k)\n",
       "                    T.reads(A[vi, vk], B[vk, vj])\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    with T.init():\n",
       "                        Y[vi, vj] = T.float32(0)\n",
       "                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "            for ax0 in T.serial(4):\n",
       "                with T.block(\"C\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + ax0)\n",
       "                    T.reads(Y[vi, vj])\n",
       "                    T.writes(C[vi, vj])\n",
       "                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_C = sch.get_block(\"C\", \"mm_relu\")\n",
    "sch.reverse_compute_at(block_C, j0)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">,</span> <span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">j_1_init</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y_init&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1_init</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">()</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y_update&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">ax0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">ax0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{for} \\PY{n}{j\\PYZus{}1\\PYZus{}init} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y\\PYZus{}init}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1\\PYZus{}init}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "            \\PY{k}{for} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y\\PYZus{}update}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                    \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "            \\PY{k}{for} \\PY{n}{ax0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{ax0}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"tir.noalias\": True, \"global_symbol\": \"mm_relu\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0 in T.grid(128, 32):\n",
       "            for j_1_init in T.serial(4):\n",
       "                with T.block(\"Y_init\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1_init)\n",
       "                    T.reads()\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "            for k, j_1 in T.grid(128, 4):\n",
       "                with T.block(\"Y_update\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                    vk = T.axis.reduce(128, k)\n",
       "                    T.reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "            for ax0 in T.serial(4):\n",
       "                with T.block(\"C\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + ax0)\n",
       "                    T.reads(Y[vi, vj])\n",
       "                    T.writes(C[vi, vj])\n",
       "                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sch.decompose_reduction(block_Y, k)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "rt_lib = tvm.build(MyModule, target=\"llvm\")\n",
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "c_nd = tvm.nd.empty((128, 128), dtype=\"float32\")\n",
    "type(c_nd)\n",
    "\n",
    "func_mm_relu = rt_lib[\"mm_relu\"]\n",
    "func_mm_relu(a_nd, b_nd, c_nd)\n",
    "\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "rt_lib = tvm.build(MyModule, target=\"llvm\")\n",
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "c_nd = tvm.nd.empty((128, 128), dtype=\"float32\")\n",
    "type(c_nd)\n",
    "\n",
    "func_mm_relu = rt_lib[\"mm_relu\"]\n",
    "func_mm_relu(a_nd, b_nd, c_nd)\n",
    "\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "rt_lib_after = tvm.build(sch.mod, target=\"llvm\")\n",
    "rt_lib_after[\"mm_relu\"](a_nd, b_nd, c_nd)\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time cost of MyModule 0.00730647 sec\n",
      "Time cost of transformed sch.mod 0.00187833 sec\n"
     ]
    }
   ],
   "source": [
    "f_timer_before = rt_lib.time_evaluator(\"mm_relu\", tvm.cpu())\n",
    "print(\"Time cost of MyModule %g sec\" % f_timer_before(a_nd, b_nd, c_nd).mean)\n",
    "f_timer_after = rt_lib_after.time_evaluator(\"mm_relu\", tvm.cpu())\n",
    "print(\"Time cost of transformed sch.mod %g sec\" % f_timer_after(a_nd, b_nd, c_nd).mean)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "要了解为什么不同的循环变体会导致不同的性能，我们需要回顾一个事实，即访问 A 和 B 中的任何内存块的速度并不一致。现代 CPU 带有多级缓存，需要先将数据提取到缓存中，然后 CPU 才能访问它。\n",
    "\n",
    "重要的是，访问已经在缓存中的数据要快得多。CPU 采用的一种策略是获取彼此更接近的数据。 当我们读取内存中的一个元素时，它会尝试将附近的元素（更正式的名称为“缓存行”）获取到缓存中。 因此，当你读取下一个元素时，它已经在缓存中。 因此，具有连续内存访问的代码通常比随机访问内存不同部分的代码更快。\n",
    "\n",
    "现在让我们看看上面的迭代可视化，分析一下是怎么回事。 在这个分析中，让我们关注最里面的两个循环：k 和 j1。高亮的地方显示了当我们针对 k 的一个特定实例迭代 j1 时迭代触及的 Y、A 和 B 中的相应区域。\n",
    "\n",
    "我们可以发现，j1 这一迭代产生了对 B 元素的连续访问。具体来说，它意味着在 j1=0 和 j1=1 时我们读取的值彼此相邻。这可以让我们拥有更好的缓存访问行为。此外，我们使 C 的计算更接近 Y，从而实现更好的缓存行为。\n",
    "\n",
    "我们当前的示例主要是为了证明不同的代码变体可以导致不同的性能。更多的变换步骤可以帮助我们获得更好的性能，我们将在以后的章节中介绍。本练习的主要目标是首先让我们获得程序变换工具，并首先体验通过变换可能实现的功能。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.4.6.1. 通过 TVMScript 创建 TensorIR¶\n",
    "获取 TensorIR 函数的第一种方法是直接在 TVMScript 中编写函数，这也是我们在上一节中使用的方法。 TVMScript 还允许我们在必要时跳过某些信息部分。 例如，T.axis.remap 使我们能够缩短迭代器大小注释。\n",
    "\n",
    "TVMScript 也是一种在变换过程中检查张量函数的有用方法。在某些情况下，一种很有帮助的做法是打印出 TVMScript，进行一些手动编辑，然后将其反馈给机器学习编译流程以调试和尝试可能的（手动）变换，然后将变换后的程序重新应用到 MLC 流程中。\n",
    "\n",
    "\n",
    "\n",
    "2.4.6.2. 使用张量表达式生成 TensorIR 代码\n",
    "在许多情况下，我们的开发形式是不在循环级别的更高级别的抽象。 所以另一种常见的获取 TensorIR 的方式是务实地生成相关代码。\n",
    "\n",
    "张量表达式 (TE) 是一种特定领域的语言，它通过 API 之类的表达式描述一系列计算。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import te\n",
    "\n",
    "A = te.placeholder((128, 128), \"float32\", name=\"A\")\n",
    "B = te.placeholder((128, 128), \"float32\", name=\"B\")\n",
    "k = te.reduce_axis((0, 128), \"k\")\n",
    "Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"Y\")\n",
    "C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name=\"C\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">,</span> <span class=\"n\">v_k</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SSR&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_k</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">v_k</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_k</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">v_k</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">v_i</span><span class=\"p\">,</span> <span class=\"n\">v_j</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{,} \\PY{n}{v\\PYZus{}k} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SSR}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}k}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{v\\PYZus{}k}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}k}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{v\\PYZus{}k}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{v\\PYZus{}i}\\PY{p}{,} \\PY{n}{v\\PYZus{}j}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j, k in T.grid(128, 128, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                v_i, v_j, v_k = T.axis.remap(\"SSR\", [i, j, k])\n",
       "                T.reads(A[v_i, v_k], B[v_k, v_j])\n",
       "                T.writes(Y[v_i, v_j])\n",
       "                with T.init():\n",
       "                    Y[v_i, v_j] = T.float32(0)\n",
       "                Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                v_i, v_j = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[v_i, v_j])\n",
       "                T.writes(C[v_i, v_j])\n",
       "                C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "te_func = te.create_prim_func([A, B, C]).with_attr({\"global_symbol\": \"mm_relu\"})\n",
    "MyModuleFromTE = tvm.IRModule({\"mm_relu\": te_func})\n",
    "IPython.display.Code(MyModuleFromTE.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mlc.ai/zh/chapter_tensor_program/tensorir_exercises.html\n",
    "2.5. TensorIR 练习\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import numpy as np\n",
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[16, 16, 16, 16],\n",
       "       [16, 16, 16, 16],\n",
       "       [16, 16, 16, 16],\n",
       "       [16, 16, 16, 16]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# init data\n",
    "a = np.arange(16).reshape(4, 4)\n",
    "b = np.arange(16, 0, -1).reshape(4, 4)\n",
    "# numpy version\n",
    "c_np = a + b\n",
    "c_np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[16, 16, 16, 16],\n",
       "       [16, 16, 16, 16],\n",
       "       [16, 16, 16, 16],\n",
       "       [16, 16, 16, 16]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# low-level numpy version\n",
    "def lnumpy_add(a: np.ndarray, b: np.ndarray, c: np.ndarray):\n",
    "  for i in range(4):\n",
    "    for j in range(4):\n",
    "      c[i, j] = a[i, j] + b[i, j]\n",
    "c_lnumpy = np.empty((4, 4), dtype=np.int64)\n",
    "lnumpy_add(a, b, c_lnumpy)\n",
    "c_lnumpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TensorIR version\n",
    "@tvm.script.ir_module\n",
    "class MyAdd:\n",
    "  @T.prim_func\n",
    "  def add(A: T.Buffer((4, 4), \"int64\"),\n",
    "          B: T.Buffer((4, 4), \"int64\"),\n",
    "          C: T.Buffer((4, 4), \"int64\")):\n",
    "    T.func_attr({\"global_symbol\": \"add\"})\n",
    "    for i, j in T.grid(4, 4):\n",
    "      with T.block(\"C\"):\n",
    "        vi = T.axis.spatial(4, i)\n",
    "        vj = T.axis.spatial(4, j)\n",
    "        C[vi, vj] = A[vi, vj] + B[vi, vj]\n",
    "\n",
    "rt_lib = tvm.build(MyAdd, target=\"llvm\")\n",
    "a_tvm = tvm.nd.array(a)\n",
    "b_tvm = tvm.nd.array(b)\n",
    "c_tvm = tvm.nd.array(np.empty((4, 4), dtype=np.int64))\n",
    "rt_lib[\"add\"](a_tvm, b_tvm, c_tvm)\n",
    "np.testing.assert_allclose(c_tvm.numpy(), c_np, rtol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.5.2. 第二节：如何变换 TensorIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">add</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">),</span> <span class=\"s2\">&quot;int64&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">),</span> <span class=\"s2\">&quot;int64&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">),</span> <span class=\"s2\">&quot;int64&quot;</span><span class=\"p\">]):</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;add&quot;</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i_0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">parallel</span><span class=\"p\">(</span><span class=\"mi\">2</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">i_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">unroll</span><span class=\"p\">(</span><span class=\"mi\">2</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">for</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">vectorized</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                    <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                        <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"n\">i_0</span> <span class=\"o\">*</span> <span class=\"mi\">2</span> <span class=\"o\">+</span> <span class=\"n\">i_1</span><span class=\"p\">)</span>\n",
       "                        <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">)</span>\n",
       "                        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                        <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{add}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{int64}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{int64}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{int64}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{add}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{k}{for} \\PY{n}{i\\PYZus{}0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{parallel}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{for} \\PY{n}{i\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{unroll}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{for} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{vectorized}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                        \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{,} \\PY{n}{i\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{2} \\PY{o}{+} \\PY{n}{i\\PYZus{}1}\\PY{p}{)}\n",
       "                        \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{,} \\PY{n}{j}\\PY{p}{)}\n",
       "                        \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                        \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                        \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def add(A: T.Buffer[(4, 4), \"int64\"], B: T.Buffer[(4, 4), \"int64\"], C: T.Buffer[(4, 4), \"int64\"]):\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"add\"})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        for i_0 in T.parallel(2):\n",
       "            for i_1 in T.unroll(2):\n",
       "                for j in T.vectorized(4):\n",
       "                    with T.block(\"C\"):\n",
       "                        vi = T.axis.spatial(4, i_0 * 2 + i_1)\n",
       "                        vj = T.axis.spatial(4, j)\n",
       "                        T.reads(A[vi, vj], B[vi, vj])\n",
       "                        T.writes(C[vi, vj])\n",
       "                        C[vi, vj] = A[vi, vj] + B[vi, vj]\n",
       "    "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyAdd:\n",
    "  @T.prim_func\n",
    "  def add(A: T.Buffer((4, 4), \"int64\"),\n",
    "          B: T.Buffer((4, 4), \"int64\"),\n",
    "          C: T.Buffer((4, 4), \"int64\")):\n",
    "    T.func_attr({\"global_symbol\": \"add\"})\n",
    "    for i, j in T.grid(4, 4):\n",
    "      with T.block(\"C\"):\n",
    "        vi = T.axis.spatial(4, i)\n",
    "        vj = T.axis.spatial(4, j)\n",
    "        C[vi, vj] = A[vi, vj] + B[vi, vj]\n",
    "\n",
    "sch = tvm.tir.Schedule(MyAdd)\n",
    "block = sch.get_block(\"C\", func_name=\"add\")\n",
    "i, j = sch.get_loops(block)\n",
    "i0, i1 = sch.split(i, factors=[2, 2])\n",
    "sch.parallel(i0)\n",
    "sch.unroll(i1)\n",
    "sch.vectorize(j)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mlc.ai/zh/chapter_end_to_end/index.html\n",
    "3. 端到端模型执行\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'relax' from 'tvm' (/home/xling/miniconda2/envs/tvm/lib/python3.9/site-packages/tvm/__init__.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[21], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mnumpy\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mnp\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtvm\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtvm\u001b[39;00m \u001b[39mimport\u001b[39;00m relax\n\u001b[1;32m      5\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtvm\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mir\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mmodule\u001b[39;00m \u001b[39mimport\u001b[39;00m IRModule\n\u001b[1;32m      6\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtvm\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mscript\u001b[39;00m \u001b[39mimport\u001b[39;00m relax \u001b[39mas\u001b[39;00m R\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'relax' from 'tvm' (/home/xling/miniconda2/envs/tvm/lib/python3.9/site-packages/tvm/__init__.py)"
     ]
    }
   ],
   "source": [
    "import IPython\n",
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "26.0%"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[23], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorchvision\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m test_data \u001b[39m=\u001b[39m torchvision\u001b[39m.\u001b[39;49mdatasets\u001b[39m.\u001b[39;49mFashionMNIST(\n\u001b[1;32m      5\u001b[0m     root\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mdata\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m      6\u001b[0m     train\u001b[39m=\u001b[39;49m\u001b[39mFalse\u001b[39;49;00m,\n\u001b[1;32m      7\u001b[0m     download\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m      8\u001b[0m     transform\u001b[39m=\u001b[39;49mtorchvision\u001b[39m.\u001b[39;49mtransforms\u001b[39m.\u001b[39;49mToTensor()\n\u001b[1;32m      9\u001b[0m )\n\u001b[1;32m     10\u001b[0m test_loader \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mutils\u001b[39m.\u001b[39mdata\u001b[39m.\u001b[39mDataLoader(test_data, batch_size\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m, shuffle\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m)\n\u001b[1;32m     11\u001b[0m class_names \u001b[39m=\u001b[39m [\u001b[39m'\u001b[39m\u001b[39mT-shirt/top\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mTrouser\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mPullover\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mDress\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mCoat\u001b[39m\u001b[39m'\u001b[39m,\n\u001b[1;32m     12\u001b[0m                \u001b[39m'\u001b[39m\u001b[39mSandal\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mShirt\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mSneaker\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mBag\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mAnkle boot\u001b[39m\u001b[39m'\u001b[39m]\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/mnist.py:99\u001b[0m, in \u001b[0;36mMNIST.__init__\u001b[0;34m(self, root, train, transform, target_transform, download)\u001b[0m\n\u001b[1;32m     96\u001b[0m     \u001b[39mreturn\u001b[39;00m\n\u001b[1;32m     98\u001b[0m \u001b[39mif\u001b[39;00m download:\n\u001b[0;32m---> 99\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdownload()\n\u001b[1;32m    101\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_check_exists():\n\u001b[1;32m    102\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39mDataset not found. You can use download=True to download it\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/mnist.py:187\u001b[0m, in \u001b[0;36mMNIST.download\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    185\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m    186\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mDownloading \u001b[39m\u001b[39m{\u001b[39;00murl\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n\u001b[0;32m--> 187\u001b[0m     download_and_extract_archive(url, download_root\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mraw_folder, filename\u001b[39m=\u001b[39;49mfilename, md5\u001b[39m=\u001b[39;49mmd5)\n\u001b[1;32m    188\u001b[0m \u001b[39mexcept\u001b[39;00m URLError \u001b[39mas\u001b[39;00m error:\n\u001b[1;32m    189\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mFailed to download (trying next):\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39m{\u001b[39;00merror\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/utils.py:434\u001b[0m, in \u001b[0;36mdownload_and_extract_archive\u001b[0;34m(url, download_root, extract_root, filename, md5, remove_finished)\u001b[0m\n\u001b[1;32m    431\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m filename:\n\u001b[1;32m    432\u001b[0m     filename \u001b[39m=\u001b[39m os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mbasename(url)\n\u001b[0;32m--> 434\u001b[0m download_url(url, download_root, filename, md5)\n\u001b[1;32m    436\u001b[0m archive \u001b[39m=\u001b[39m os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mjoin(download_root, filename)\n\u001b[1;32m    437\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mExtracting \u001b[39m\u001b[39m{\u001b[39;00marchive\u001b[39m}\u001b[39;00m\u001b[39m to \u001b[39m\u001b[39m{\u001b[39;00mextract_root\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/utils.py:144\u001b[0m, in \u001b[0;36mdownload_url\u001b[0;34m(url, root, filename, md5, max_redirect_hops)\u001b[0m\n\u001b[1;32m    142\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m    143\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mDownloading \u001b[39m\u001b[39m\"\u001b[39m \u001b[39m+\u001b[39m url \u001b[39m+\u001b[39m \u001b[39m\"\u001b[39m\u001b[39m to \u001b[39m\u001b[39m\"\u001b[39m \u001b[39m+\u001b[39m fpath)\n\u001b[0;32m--> 144\u001b[0m     _urlretrieve(url, fpath)\n\u001b[1;32m    145\u001b[0m \u001b[39mexcept\u001b[39;00m (urllib\u001b[39m.\u001b[39merror\u001b[39m.\u001b[39mURLError, \u001b[39mOSError\u001b[39;00m) \u001b[39mas\u001b[39;00m e:  \u001b[39m# type: ignore[attr-defined]\u001b[39;00m\n\u001b[1;32m    146\u001b[0m     \u001b[39mif\u001b[39;00m url[:\u001b[39m5\u001b[39m] \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mhttps\u001b[39m\u001b[39m\"\u001b[39m:\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/utils.py:48\u001b[0m, in \u001b[0;36m_urlretrieve\u001b[0;34m(url, filename, chunk_size)\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_urlretrieve\u001b[39m(url: \u001b[39mstr\u001b[39m, filename: \u001b[39mstr\u001b[39m, chunk_size: \u001b[39mint\u001b[39m \u001b[39m=\u001b[39m \u001b[39m1024\u001b[39m \u001b[39m*\u001b[39m \u001b[39m32\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     47\u001b[0m     \u001b[39mwith\u001b[39;00m urllib\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39murlopen(urllib\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39mRequest(url, headers\u001b[39m=\u001b[39m{\u001b[39m\"\u001b[39m\u001b[39mUser-Agent\u001b[39m\u001b[39m\"\u001b[39m: USER_AGENT})) \u001b[39mas\u001b[39;00m response:\n\u001b[0;32m---> 48\u001b[0m         _save_response_content(\u001b[39miter\u001b[39;49m(\u001b[39mlambda\u001b[39;49;00m: response\u001b[39m.\u001b[39;49mread(chunk_size), \u001b[39mb\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39m\"\u001b[39;49m), filename, length\u001b[39m=\u001b[39;49mresponse\u001b[39m.\u001b[39;49mlength)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/utils.py:37\u001b[0m, in \u001b[0;36m_save_response_content\u001b[0;34m(content, destination, length)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_save_response_content\u001b[39m(\n\u001b[1;32m     32\u001b[0m     content: Iterator[\u001b[39mbytes\u001b[39m],\n\u001b[1;32m     33\u001b[0m     destination: \u001b[39mstr\u001b[39m,\n\u001b[1;32m     34\u001b[0m     length: Optional[\u001b[39mint\u001b[39m] \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m     35\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     36\u001b[0m     \u001b[39mwith\u001b[39;00m \u001b[39mopen\u001b[39m(destination, \u001b[39m\"\u001b[39m\u001b[39mwb\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39mas\u001b[39;00m fh, tqdm(total\u001b[39m=\u001b[39mlength) \u001b[39mas\u001b[39;00m pbar:\n\u001b[0;32m---> 37\u001b[0m         \u001b[39mfor\u001b[39;00m chunk \u001b[39min\u001b[39;00m content:\n\u001b[1;32m     38\u001b[0m             \u001b[39m# filter out keep-alive new chunks\u001b[39;00m\n\u001b[1;32m     39\u001b[0m             \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m chunk:\n\u001b[1;32m     40\u001b[0m                 \u001b[39mcontinue\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/torchvision/datasets/utils.py:48\u001b[0m, in \u001b[0;36m_urlretrieve.<locals>.<lambda>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_urlretrieve\u001b[39m(url: \u001b[39mstr\u001b[39m, filename: \u001b[39mstr\u001b[39m, chunk_size: \u001b[39mint\u001b[39m \u001b[39m=\u001b[39m \u001b[39m1024\u001b[39m \u001b[39m*\u001b[39m \u001b[39m32\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     47\u001b[0m     \u001b[39mwith\u001b[39;00m urllib\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39murlopen(urllib\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39mRequest(url, headers\u001b[39m=\u001b[39m{\u001b[39m\"\u001b[39m\u001b[39mUser-Agent\u001b[39m\u001b[39m\"\u001b[39m: USER_AGENT})) \u001b[39mas\u001b[39;00m response:\n\u001b[0;32m---> 48\u001b[0m         _save_response_content(\u001b[39miter\u001b[39m(\u001b[39mlambda\u001b[39;00m: response\u001b[39m.\u001b[39;49mread(chunk_size), \u001b[39mb\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m\"\u001b[39m), filename, length\u001b[39m=\u001b[39mresponse\u001b[39m.\u001b[39mlength)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/http/client.py:459\u001b[0m, in \u001b[0;36mHTTPResponse.read\u001b[0;34m(self, amt)\u001b[0m\n\u001b[1;32m    456\u001b[0m \u001b[39mif\u001b[39;00m amt \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    457\u001b[0m     \u001b[39m# Amount is given, implement using readinto\u001b[39;00m\n\u001b[1;32m    458\u001b[0m     b \u001b[39m=\u001b[39m \u001b[39mbytearray\u001b[39m(amt)\n\u001b[0;32m--> 459\u001b[0m     n \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mreadinto(b)\n\u001b[1;32m    460\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mmemoryview\u001b[39m(b)[:n]\u001b[39m.\u001b[39mtobytes()\n\u001b[1;32m    461\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    462\u001b[0m     \u001b[39m# Amount is not given (unbounded read) so we must check self.length\u001b[39;00m\n\u001b[1;32m    463\u001b[0m     \u001b[39m# and self.chunked\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/http/client.py:503\u001b[0m, in \u001b[0;36mHTTPResponse.readinto\u001b[0;34m(self, b)\u001b[0m\n\u001b[1;32m    498\u001b[0m         b \u001b[39m=\u001b[39m \u001b[39mmemoryview\u001b[39m(b)[\u001b[39m0\u001b[39m:\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlength]\n\u001b[1;32m    500\u001b[0m \u001b[39m# we do not use _safe_read() here because this may be a .will_close\u001b[39;00m\n\u001b[1;32m    501\u001b[0m \u001b[39m# connection, and the user is reading more bytes than will be provided\u001b[39;00m\n\u001b[1;32m    502\u001b[0m \u001b[39m# (for example, reading in 1k chunks)\u001b[39;00m\n\u001b[0;32m--> 503\u001b[0m n \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mfp\u001b[39m.\u001b[39;49mreadinto(b)\n\u001b[1;32m    504\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m n \u001b[39mand\u001b[39;00m b:\n\u001b[1;32m    505\u001b[0m     \u001b[39m# Ideally, we would raise IncompleteRead if the content-length\u001b[39;00m\n\u001b[1;32m    506\u001b[0m     \u001b[39m# wasn't satisfied, but it might break compatibility.\u001b[39;00m\n\u001b[1;32m    507\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_close_conn()\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/socket.py:669\u001b[0m, in \u001b[0;36mSocketIO.readinto\u001b[0;34m(self, b)\u001b[0m\n\u001b[1;32m    667\u001b[0m \u001b[39mwhile\u001b[39;00m \u001b[39mTrue\u001b[39;00m:\n\u001b[1;32m    668\u001b[0m     \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 669\u001b[0m         \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_sock\u001b[39m.\u001b[39;49mrecv_into(b)\n\u001b[1;32m    670\u001b[0m     \u001b[39mexcept\u001b[39;00m timeout:\n\u001b[1;32m    671\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_timeout_occurred \u001b[39m=\u001b[39m \u001b[39mTrue\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "test_data = torchvision.datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor()\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)\n",
    "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
    "               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
    "\n",
    "img, label = next(iter(test_loader))\n",
    "img = img.reshape(1, 28, 28).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img[0])\n",
    "plt.colorbar()\n",
    "plt.grid(False)\n",
    "plt.show()\n",
    "print(\"Class:\", class_names[label[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "# Hide outputs\n",
    "!wget https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_params.pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def numpy_mlp(data, w0, b0, w1, b1):\n",
    "    lv0 = data @ w0.T + b0\n",
    "    lv1 = np.maximum(lv0, 0)\n",
    "    lv2 = lv1 @ w1.T + b1\n",
    "    return lv2\n",
    "\n",
    "import pickle as pkl\n",
    "\n",
    "mlp_params = pkl.load(open(\"fasionmnist_mlp_params.pkl\", \"rb\"))\n",
    "res = numpy_mlp(img.reshape(1, 784),\n",
    "                mlp_params[\"w0\"],\n",
    "                mlp_params[\"b0\"],\n",
    "                mlp_params[\"w1\"],\n",
    "                mlp_params[\"b1\"])\n",
    "print(res)\n",
    "pred_kind = res.argmax(axis=1)\n",
    "print(pred_kind)\n",
    "print(\"NumPy-MLP Prediction:\", class_names[pred_kind[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3.4. 在 TVMScript 中构建端到端 IRModule¶\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def relu0(x: T.handle, y: T.handle):\n",
    "        n = T.int64()\n",
    "        X = T.match_buffer(x, (1, n), \"float32\")\n",
    "        Y = T.match_buffer(y, (1, n), \"float32\")\n",
    "        for i, j in T.grid(1, n):\n",
    "            with T.block(\"Y\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))\n",
    "\n",
    "    @T.prim_func\n",
    "    def linear0(x: T.handle,\n",
    "                w: T.handle,\n",
    "                b: T.handle,\n",
    "                z: T.handle):\n",
    "        m, n, k = T.int64(), T.int64(), T.int64()\n",
    "        X = T.match_buffer(x, (1, m), \"float32\")\n",
    "        W = T.match_buffer(w, (n, m), \"float32\")\n",
    "        B = T.match_buffer(b, (n, ), \"float32\")\n",
    "        Z = T.match_buffer(z, (1, n), \"float32\")\n",
    "        Y = T.alloc_buffer((1, n), \"float32\")\n",
    "        for i, j, k in T.grid(1, n, m):\n",
    "            with T.block(\"Y\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
    "        for i, j in T.grid(1, n):\n",
    "            with T.block(\"Z\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                Z[vi, vj] = Y[vi, vj] + B[vj]\n",
    "\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((1, \"m\"), \"float32\"),\n",
    "             w0: R.Tensor((\"n\", \"m\"), \"float32\"),\n",
    "             b0: R.Tensor((\"n\", ), \"float32\"),\n",
    "             w1: R.Tensor((\"k\", \"n\"), \"float32\"),\n",
    "             b1: R.Tensor((\"k\", ), \"float32\")):\n",
    "        m, n, k = T.int64(), T.int64(), T.int64()\n",
    "        with R.dataflow():\n",
    "            lv0 = R.call_dps_packed(\"linear0\", (x, w0, b0), R.Tensor((1, n), \"float32\"))\n",
    "            lv1 = R.call_dps_packed(\"relu0\", (lv0, ), R.Tensor((1, n), \"float32\"))\n",
    "            out = R.call_dps_packed(\"linear0\", (lv1, w1, b1), R.Tensor((1, k), \"float32\"))\n",
    "            R.output(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ex = relax.build(MyModule, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "data_nd = tvm.nd.array(img.reshape(1, 784))\n",
    "nd_params = {k: tvm.nd.array(v) for k, v in mlp_params.items()}\n",
    "nd_res = vm[\"main\"](data_nd,\n",
    "                    nd_params[\"w0\"],\n",
    "                    nd_params[\"b0\"],\n",
    "                    nd_params[\"w1\"],\n",
    "                    nd_params[\"b1\"])\n",
    "print(nd_res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
    "print(\"MyModule Prediction:\", class_names[pred_kind[0]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3.6. 在环境中集成现有运行库\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3.6.1. 注册运行时函数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.register_func(\"env.linear\", override=True)\n",
    "def torch_linear(x: tvm.nd.NDArray,\n",
    "                 w: tvm.nd.NDArray,\n",
    "                 b: tvm.nd.NDArray,\n",
    "                 out: tvm.nd.NDArray):\n",
    "    x_torch = torch.from_dlpack(x)\n",
    "    w_torch = torch.from_dlpack(w)\n",
    "    b_torch = torch.from_dlpack(b)\n",
    "    out_torch = torch.from_dlpack(out)\n",
    "    torch.mm(x_torch, w_torch.T, out=out_torch)\n",
    "    torch.add(out_torch, b_torch, out=out_torch)\n",
    "\n",
    "@tvm.register_func(\"env.relu\", override=True)\n",
    "def lnumpy_relu(x: tvm.nd.NDArray,\n",
    "                out: tvm.nd.NDArray):\n",
    "    x_torch = torch.from_dlpack(x)\n",
    "    out_torch = torch.from_dlpack(out)\n",
    "    torch.maximum(x_torch, torch.Tensor([0.0]), out=out_torch)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModuleWithExternCall:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((1, \"m\"), \"float32\"),\n",
    "             w0: R.Tensor((\"n\", \"m\"), \"float32\"),\n",
    "             b0: R.Tensor((\"n\", ), \"float32\"),\n",
    "             w1: R.Tensor((\"k\", \"n\"), \"float32\"),\n",
    "             b1: R.Tensor((\"k\", ), \"float32\")):\n",
    "        # block 0\n",
    "        m, n, k = T.int64(), T.int64(), T.int64()\n",
    "        with R.dataflow():\n",
    "            lv0 = R.call_dps_packed(\"env.linear\", (x, w0, b0), R.Tensor((1, n), \"float32\"))\n",
    "            lv1 = R.call_dps_packed(\"env.relu\", (lv0, ), R.Tensor((1, n), \"float32\"))\n",
    "            out = R.call_dps_packed(\"env.linear\", (lv1, w1, b1), R.Tensor((1, k), \"float32\"))\n",
    "            R.output(out)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ex = relax.build(MyModuleWithExternCall, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "\n",
    "nd_res = vm[\"main\"](data_nd,\n",
    "                    nd_params[\"w0\"],\n",
    "                    nd_params[\"b0\"],\n",
    "                    nd_params[\"w1\"],\n",
    "                    nd_params[\"b1\"])\n",
    "\n",
    "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
    "print(\"MyModuleWithExternCall Prediction:\", class_names[pred_kind[0]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. 自动程序优化¶\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import IPython\n",
    "\n",
    "\n",
    "def code2html(code):\n",
    "    \"\"\"Helper function to use pygments to turn the code string into highlighted html.\"\"\"\n",
    "    import pygments\n",
    "    from pygments.formatters import HtmlFormatter\n",
    "    from pygments.lexers import Python3Lexer\n",
    "    formatter = HtmlFormatter()\n",
    "    html = pygments.highlight(code, Python3Lexer(), formatter)\n",
    "    return \"<style>%s</style>%s\\n\" % (formatter.get_style_defs(\".highlight\"), html)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def main(\n",
    "        A: T.Buffer((128, 128), \"float32\"),\n",
    "        B: T.Buffer((128, 128), \"float32\"),\n",
    "        C: T.Buffer((128, 128), \"float32\"),\n",
    "    ):\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    C[vi, vj] = 0.0\n",
    "                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "dtype = \"float32\"\n",
    "a_np = np.random.rand(128, 128).astype(dtype)\n",
    "b_np = np.random.rand(128, 128).astype(dtype)\n",
    "c_mm = a_np @ b_np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "c_nd = tvm.nd.empty((128, 128), dtype=\"float32\")\n",
    "\n",
    "lib = tvm.build(MyModule, target=\"llvm\")\n",
    "f_timer_before = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "print(\"Time cost of MyModule: %.3f ms\" % (f_timer_before(a_nd, b_nd, c_nd).mean * 1000))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def schedule_mm(sch: tvm.tir.Schedule, jfactor=4):\n",
    "    block_C = sch.get_block(\"C\", \"main\")\n",
    "    i, j, k = sch.get_loops(block=block_C)\n",
    "    j_0, j_1 = sch.split(loop=j, factors=[None, jfactor])\n",
    "    sch.reorder(i, j_0, k, j_1)\n",
    "    sch.decompose_reduction(block_C, k)\n",
    "    return sch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(MyModule)\n",
    "sch = schedule_mm(sch)\n",
    "IPython.display.HTML(code2html(sch.mod.script()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "print(\"Time cost of MyModule=>schedule_mm: %.3f ms\" % (f_timer_after(a_nd, b_nd, c_nd).mean * 1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "print(sch.trace)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4.4. 随机调度变换 (Stochastic Schedule Transformation)¶\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def stochastic_schedule_mm(sch: tvm.tir.Schedule):\n",
    "    block_C = sch.get_block(\"C\", \"main\")\n",
    "    i, j, k = sch.get_loops(block=block_C)\n",
    "    j_factors = sch.sample_perfect_tile(loop=j, n=2)\n",
    "    j_0, j_1 = sch.split(loop=j, factors=j_factors)\n",
    "    sch.reorder(i, j_0, k, j_1)\n",
    "    sch.decompose_reduction(block_C, k)\n",
    "    return sch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4.5. 随机变换搜索\n",
    "你可能会意识到，stochastic_schedule_mm 创建了一个可能程序的搜索空间，具体取决于在每个采样步骤中做出的具体决定。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def random_search(mod: tvm.IRModule, num_trials=5):\n",
    "    best_result = None\n",
    "    best_sch = None\n",
    "\n",
    "    for i in range(num_trials):\n",
    "        sch = stochastic_schedule_mm(tvm.tir.Schedule(mod))\n",
    "        lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "        f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "        result = f_timer_after(a_nd, b_nd, c_nd).mean\n",
    "\n",
    "        print(\"=====Attempt %d, time-cost: %.3f ms====\" % (i, result * 1000))\n",
    "        print(sch.trace)\n",
    "\n",
    "        # book keep the best result so far\n",
    "        if best_result is None or result < best_result:\n",
    "            best_result = result\n",
    "            best_sch = sch\n",
    "\n",
    "    return best_sch\n",
    "\n",
    "sch = random_search(MyModule)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "print(sch.trace)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "from tvm import meta_schedule as ms\n",
    "\n",
    "database = ms.tune_tir(\n",
    "    mod=MyModule,\n",
    "    target=\"llvm --num-cores=1\",\n",
    "    max_trials_global=64,\n",
    "    num_trials_per_iter=64,\n",
    "    space=ms.space_generator.ScheduleFn(stochastic_schedule_mm),\n",
    "    work_dir=\"./tune_tmp\",\n",
    "    # task_name=\"main\"\n",
    ")\n",
    "\n",
    "sch = ms.tir_integration.compile_tir(database, MyModule, \"llvm --num-cores=1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "print(\"Time cost of MyModule after tuning: %.3f ms\" % (f_timer_after(a_nd, b_nd, c_nd).mean * 1000))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4.6. 回到端到端模型执行\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "test_data = torchvision.datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor()\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)\n",
    "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
    "               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
    "\n",
    "img, label = next(iter(test_loader))\n",
    "img = img.reshape(1, 28, 28).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import pickle as pkl\n",
    "\n",
    "mlp_params = pkl.load(open(\"fasionmnist_mlp_params.pkl\", \"rb\"))\n",
    "\n",
    "data_nd = tvm.nd.array(img.reshape(1, 784))\n",
    "nd_params = {k: tvm.nd.array(v) for k, v in mlp_params.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModuleMixture:\n",
    "    @T.prim_func\n",
    "    def linear0(X: T.Buffer((1, 784), \"float32\"),\n",
    "                W: T.Buffer((128, 784), \"float32\"),\n",
    "                B: T.Buffer((128,), \"float32\"),\n",
    "                Z: T.Buffer((1, 128), \"float32\")):\n",
    "        T.func_attr({\"global_symbol\": \"linear0\", \"tir.noalias\": True})\n",
    "        Y = T.alloc_buffer((1, 128), \"float32\")\n",
    "        for i, j, k in T.grid(1, 128, 784):\n",
    "            with T.block(\"Y\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
    "\n",
    "        for i, j in T.grid(1, 128):\n",
    "            with T.block(\"Z\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                Z[vi, vj] =  Y[vi, vj] + B[vj]\n",
    "\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((1, 784), \"float32\"),\n",
    "             w0: R.Tensor((128, 784), \"float32\"),\n",
    "             b0: R.Tensor((128,), \"float32\"),\n",
    "             w1: R.Tensor((10, 128), \"float32\"),\n",
    "             b1: R.Tensor((10,), \"float32\")):\n",
    "        with R.dataflow():\n",
    "            lv0 = R.call_dps_packed(\"linear0\", (x, w0, b0), R.Tensor((1, 128), dtype=\"float32\"))\n",
    "            lv1 = R.call_dps_packed(\"env.relu\", (lv0,), R.Tensor((1, 128), dtype=\"float32\"))\n",
    "            out = R.call_dps_packed(\"env.linear\", (lv1, w1, b1), R.Tensor((1, 10), dtype=\"float32\"))\n",
    "            R.output(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "@tvm.register_func(\"env.linear\", override=True)\n",
    "def torch_linear(x: tvm.nd.NDArray,\n",
    "                 w: tvm.nd.NDArray,\n",
    "                 b: tvm.nd.NDArray,\n",
    "                 out: tvm.nd.NDArray):\n",
    "    x_torch = torch.from_dlpack(x)\n",
    "    w_torch = torch.from_dlpack(w)\n",
    "    b_torch = torch.from_dlpack(b)\n",
    "    out_torch = torch.from_dlpack(out)\n",
    "    torch.mm(x_torch, w_torch.T, out=out_torch)\n",
    "    torch.add(out_torch, b_torch, out=out_torch)\n",
    "\n",
    "@tvm.register_func(\"env.relu\", override=True)\n",
    "def lnumpy_relu(x: tvm.nd.NDArray,\n",
    "                out: tvm.nd.NDArray):\n",
    "    x_torch = torch.from_dlpack(x)\n",
    "    out_torch = torch.from_dlpack(out)\n",
    "    torch.maximum(x_torch, torch.Tensor([0.0]), out=out_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "MyModuleWithParams = relax.transform.BindParams(\"main\", nd_params)(MyModuleMixture)\n",
    "ex = relax.build(MyModuleWithParams, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "\n",
    "nd_res = vm[\"main\"](data_nd)\n",
    "\n",
    "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
    "print(\"MyModuleWithParams Prediction:\", class_names[pred_kind[0]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ftimer = vm.module.time_evaluator(\"main\", tvm.cpu(), number=100)\n",
    "\n",
    "print(\"MyModuleWithParams time-cost: %g ms\" % (ftimer(data_nd).mean * 1000))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "目前，调优 API 只接受一个带有一个 main 函数的 IRModule，所以我们首先将 linear0 取出到另一个模块的 main 函数中并将其传递给 tune_tir。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "mod_linear = tvm.IRModule.from_expr(MyModuleMixture[\"linear0\"].with_attr(\"global_symbol\", \"main\"))\n",
    "IPython.display.HTML(code2html(mod_linear.script()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "database = ms.tune_tir(\n",
    "    mod=mod_linear,\n",
    "    target=\"llvm --num-cores=1\",\n",
    "    max_trials_global=64,\n",
    "    num_trials_per_iter=64,\n",
    "    work_dir=\"./tune_tmp\",\n",
    "    # task_name=\"main\",\n",
    ")\n",
    "sch = ms.tir_integration.compile_tir(database, mod_linear, \"llvm --num-cores=1\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "MyModuleWithParams2 = relax.transform.BindParams(\"main\", nd_params)(MyModuleMixture)\n",
    "new_func = sch.mod[\"main\"].with_attr(\"global_symbol\", \"linear0\")\n",
    "gv = MyModuleWithParams2.get_global_var(\"linear0\")\n",
    "MyModuleWithParams2.update_func(gv, new_func)\n",
    "IPython.display.HTML(code2html(MyModuleWithParams2.script()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ex = relax.build(MyModuleWithParams2, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "\n",
    "nd_res = vm[\"main\"](data_nd)\n",
    "\n",
    "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
    "print(\"MyModuleWithParams2 Prediction:\", class_names[pred_kind[0]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ftimer = vm.module.time_evaluator(\"main\", tvm.cpu(), number=50)\n",
    "\n",
    "print(\"MyModuleWithParams2 time-cost: %g ms\" % (ftimer(data_nd).mean * 1000))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. 与机器学习框架的整合\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import fx\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "from tvm import te"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "A = te.placeholder((128, 128), name=\"A\", dtype=\"float32\")\n",
    "B = te.placeholder((128, 128), name=\"B\", dtype=\"float32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def te_matmul(A: te.Tensor, B: te.Tensor) -> te.Tensor:\n",
    "    assert A.shape[1] == B.shape[0]\n",
    "    n = A.shape[0]\n",
    "    m = B.shape[1]\n",
    "    k = te.reduce_axis((0, A.shape[1]), name=\"k\")\n",
    "    return te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"matmul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "C = te_matmul(A, B)\n",
    "te.create_prim_func([A, B, C]).show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def te_relu(A: te.Tensor) -> te.Tensor:\n",
    "    return te.compute(A.shape, lambda *i: te.max(A(*i), 0), name=\"relu\")\n",
    "\n",
    "X1 = te.placeholder((10,), name=\"X1\", dtype=\"float32\")\n",
    "Y1 = te_relu(X1)\n",
    "te.create_prim_func([X1, Y1]).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "X2 = te.placeholder((10, 20), name=\"X1\", dtype=\"float32\")\n",
    "Y2 = te_relu(X2)\n",
    "te.create_prim_func([X2, Y2]).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "te API 允许我们做的另一件事是组合操作并创建“融合 (fused)”算子。例如，我们可以将 matmul 的结果再次应用 relu。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "C = te_matmul(A, B)\n",
    "D = te_relu(C)\n",
    "te.create_prim_func([A, B, D]).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.3.2. 使用 BlockBuilder 构造 IRModule¶\n",
    "到目前为止，我们已经创建了一个 TensorIR 函数。 为了构建端到端的模型执行，我们还需要能够通过计算图连接多个 TensorIR 函数。\n",
    "\n",
    "让我们首先创建一个 block builder，它可以帮助我们逐步构建一个 relax.Function。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "A = relax.Var(\"A\", relax.TensorStructInfo((128, 128), \"float32\"))\n",
    "B = relax.Var(\"B\", relax.TensorStructInfo((128, 128), \"float32\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "bb = relax.BlockBuilder()\n",
    "\n",
    "with bb.function(\"main\"):\n",
    "    with bb.dataflow():\n",
    "        C = bb.emit_te(te_matmul, A, B)\n",
    "        D = bb.emit_te(te_relu, C)\n",
    "        R = bb.emit_output(D)\n",
    "    bb.emit_func_output(R, params=[A, B])\n",
    "\n",
    "MyModule = bb.get()\n",
    "MyModule.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.3.3. 深入理解 BlockBuilder API¶\n",
    "现在让我们深入了解 BlockBuilder 的 API。将 BlockBuilder 代码和生成的 IRModule 并排放置会很有帮助。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.4. 从 PyTorch 导入模型¶\n",
    "现在我们已经学习了以编程方式构建 IRModule 的工具。 让我们使用它们将机器学习模型从 PyTorch 导入成为 IRModule。\n",
    "\n",
    "大多数机器学习框架都带有计算图抽象，其中每个节点对应一个操作，边对应它们之间的依赖关系。 我们将采用 PyTorch 模型，获取 PyTorch 原生格式的计算图，并将其转换为 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "class MyModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.weight = nn.Parameter(torch.randn(128, 128))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.matmul(x, self.weight)\n",
    "        x = torch.relu(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "model = MyModel()\n",
    "fx_module = fx.symbolic_trace(model)\n",
    "type(fx_module)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "fx_module.graph.print_tabular()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.4.2. 构造映射函数¶\n",
    "让我们定义整体的翻译逻辑。 主要流程如下：\n",
    "\n",
    "创建一个 node_map，将 fx.Node 映射到相应的 relax.Var，该 relax.Var 代表 IRModule 中的已翻译节点。\n",
    "\n",
    "以拓扑顺序迭代 FX 图中的节点。\n",
    "\n",
    "给定映射输入，获取节点的映射输出。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def map_param(param: nn.Parameter):\n",
    "    return relax.const(\n",
    "        param.data.cpu().numpy(), relax.TensorStructInfo(param.data.shape, \"float32\")\n",
    "    )\n",
    "\n",
    "def fetch_attr(fx_mod, target: str):\n",
    "    \"\"\"Helper function to fetch an attr\"\"\"\n",
    "    target_atoms = target.split('.')\n",
    "    attr_itr = fx_mod\n",
    "    for i, atom in enumerate(target_atoms):\n",
    "        if not hasattr(attr_itr, atom):\n",
    "            raise RuntimeError(f\"Node referenced nonexistant target {'.'.join(target_atoms[:i])}\")\n",
    "        attr_itr = getattr(attr_itr, atom)\n",
    "    return attr_itr\n",
    "\n",
    "def from_fx(fx_mod, input_shapes, call_function_map, call_module_map):\n",
    "    input_index = 0\n",
    "    node_map = {}\n",
    "    named_modules = dict(fx_mod.named_modules())\n",
    "\n",
    "    bb = relax.BlockBuilder()\n",
    "\n",
    "    fn_inputs = []\n",
    "    fn_output = None\n",
    "    with bb.function(\"main\"):\n",
    "        with bb.dataflow():\n",
    "            for node in fx_mod.graph.nodes:\n",
    "                if node.op == \"placeholder\":\n",
    "                    # create input placeholder\n",
    "                    shape = input_shapes[input_index]\n",
    "                    input_index += 1\n",
    "                    input_var = relax.Var(\n",
    "                        node.target, relax.TensorStructInfo(shape, \"float32\")\n",
    "                    )\n",
    "                    fn_inputs.append(input_var)\n",
    "                    node_map[node] = input_var\n",
    "                elif node.op == \"get_attr\":\n",
    "                    node_map[node] = map_param(fetch_attr(fx_mod, node.target))\n",
    "                elif node.op == \"call_function\":\n",
    "                    node_map[node] = call_function_map[node.target](bb, node_map, node)\n",
    "                elif node.op == \"call_module\":\n",
    "                    named_module = named_modules[node.target]\n",
    "                    node_map[node] = call_module_map[type(named_module)](bb, node_map, node, named_module)\n",
    "                elif node.op == \"output\":\n",
    "                    output = node_map[node.args[0]]\n",
    "                    assert fn_output is None\n",
    "                    fn_output = bb.emit_output(output)\n",
    "        # output and finalize the function\n",
    "        bb.emit_func_output(output, fn_inputs)\n",
    "    return bb.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "def map_matmul(bb, node_map, node: fx.Node):\n",
    "    A = node_map[node.args[0]]\n",
    "    B = node_map[node.args[1]]\n",
    "    return bb.emit_te(te_matmul, A, B)\n",
    "\n",
    "def map_relu(bb, node_map, node: fx.Node):\n",
    "    A = node_map[node.args[0]]\n",
    "    return bb.emit_te(te_relu, A)\n",
    "\n",
    "MyModule = from_fx(\n",
    "    fx_module,\n",
    "    input_shapes = [(1, 128)],\n",
    "    call_function_map = {\n",
    "      torch.matmul: map_matmul,\n",
    "      torch.relu: map_relu,\n",
    "    },\n",
    "    call_module_map={},\n",
    ")\n",
    "\n",
    "MyModule.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.5. 回到 FashionMNIST 的例子\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "test_data = torchvision.datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor()\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)\n",
    "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
    "               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
    "\n",
    "img, label = next(iter(test_loader))\n",
    "img = img.reshape(1, 28, 28).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        self.linear0 = nn.Linear(784, 128, bias=True)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.linear1 = nn.Linear(128, 10, bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear0(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.linear1(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import pickle as pkl\n",
    "\n",
    "mlp_model = MLP()\n",
    "\n",
    "mlp_params = pkl.load(open(\"fasionmnist_mlp_params.pkl\", \"rb\"))\n",
    "mlp_model.linear0.weight.data = torch.from_numpy(mlp_params[\"w0\"])\n",
    "mlp_model.linear0.bias.data = torch.from_numpy(mlp_params[\"b0\"])\n",
    "mlp_model.linear1.weight.data = torch.from_numpy(mlp_params[\"w1\"])\n",
    "mlp_model.linear1.bias.data = torch.from_numpy(mlp_params[\"b1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "torch_res = mlp_model(torch.from_numpy(img.reshape(1, 784)))\n",
    "\n",
    "pred_kind = np.argmax(torch_res.detach().numpy(), axis=1)\n",
    "print(\"Torch Prediction:\", class_names[pred_kind[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "from tvm import topi\n",
    "\n",
    "\n",
    "def map_nn_linear(bb, node_map, node, nn_mod):\n",
    "    x = node_map[node.args[0]]\n",
    "    w = map_param(nn_mod.weight)\n",
    "    if nn_mod.bias is not None:\n",
    "        b = map_param(nn_mod.bias)\n",
    "    y = bb.emit_te(topi.nn.dense, x, w)\n",
    "    return bb.emit_te(topi.add, y, b)\n",
    "\n",
    "def map_nn_relu(bb, node_map, node, nn_mod):\n",
    "    return map_relu(bb, node_map, node)\n",
    "\n",
    "\n",
    "MLPModule = from_fx(\n",
    "    fx.symbolic_trace(mlp_model),\n",
    "    input_shapes = [(1, 784)],\n",
    "    call_function_map={\n",
    "    },\n",
    "    call_module_map={\n",
    "        torch.nn.Linear: map_nn_linear,\n",
    "        torch.nn.ReLU: map_nn_relu,\n",
    "    },\n",
    ")\n",
    "\n",
    "MLPModule.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'base' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'conda install -n base ipykernel --update-deps --force-reinstall'"
     ]
    }
   ],
   "source": [
    "ex = relax.build(MLPModule, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "data_nd = tvm.nd.array(img.reshape(1, 784))\n",
    "\n",
    "nd_res = vm[\"main\"](data_nd)\n",
    "\n",
    "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
    "print(\"MLPModule Prediction:\", class_names[pred_kind[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1. 第一部分\n",
    "在上一章中，我们讨论了 CPU 环境中的 MLC 流程。本章节将讨论如何将一些优化带到 GPU 上。我们将使用 CUDA 编程语言。然而，同样的概念也适用于其他类型的 GPU。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModuleVecAdd:\n",
    "    @T.prim_func\n",
    "    def main(A: T.Buffer((1024,), \"float32\"),\n",
    "             B: T.Buffer((1024,), \"float32\"),\n",
    "             C: T.Buffer((1024,), \"float32\")) -> None:\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i in T.grid(1024):\n",
    "            with T.block(\"C\"):\n",
    "                vi = T.axis.remap(\"S\", [i])\n",
    "                C[vi] = A[vi] + B[vi]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xling/miniconda2/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/home/xling/miniconda2/envs/mlc/bin/python -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0, i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">128</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi], B[vi])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi])\n",
       "                C[vi] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[vi] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[vi]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(MyModuleVecAdd)\n",
    "block_C = sch.get_block(\"C\")\n",
    "i, = sch.get_loops(block=block_C)\n",
    "i0, i1 = sch.split(i, [None, 128])\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "TVMError",
     "evalue": "Traceback (most recent call last):\n  12: TVMFuncCall\n  11: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::Module (tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)>::AssignTypedLambda<tvm::__mk_TVM22::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#1}>(tvm::__mk_TVM22::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tvm::runtime::TVMRetValue)\n  10: tvm::TIRToRuntime(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target const&)\n  9: tvm::SplitMixedModule(tvm::IRModule, tvm::Target const&, tvm::Target const&)\n  8: tvm::ApplyPasses(tvm::IRModule, tvm::transform::Sequential)\n  7: tvm::transform::Pass::operator()(tvm::IRModule) const\n  6: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  5: tvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  4: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  3: tvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  2: _ZN3tvm7runtime13PackedFun\n  1: tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::tir::transform::VerifyMemory()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1}>(tvm::tir::transform::VerifyMemory()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1})::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  0: _ZN3tvm7runtime6deta\n  Did you forget to bind?\n    Variable `B` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n    Variable `A` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n    Variable `C` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n  File \"/workspace/tvm/src/tir/analysis/verify_memory.cc\", line 205\nRuntimeError: Memory verification failed with the following errors:\n# from tvm.script import tir as T\n\n@T.prim_func\ndef main(A: T.Buffer((1024,), \"float32\"), B: T.Buffer((1024,), \"float32\"), C: T.Buffer((1024,), \"float32\")):\n    T.func_attr({\"target\": T.target({\"arch\": \"sm_75\", \"host\": {\"keys\": [\"cpu\"], \"kind\": \"llvm\", \"tag\": \"\"}, \"keys\": [\"cuda\", \"gpu\"], \"kind\": \"cuda\", \"max_num_threads\": 1024, \"tag\": \"\", \"thread_warp_size\": 32}), \"tir.noalias\": T.bool(True)})\n    for i_0, i_1 in T.grid(8, 128):\n        cse_var_1: T.int32 = i_0 * 128 + i_1\n        C_1 = T.Buffer((1024,), data=C.data)\n        A_1 = T.Buffer((1024,), data=A.data)\n        B_1 = T.Buffer((1024,), data=B.data)\n        C_1[cse_var_1] = A_1[cse_var_1] + B_1[cse_var_1]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTVMError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[29], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m rt_mod \u001b[39m=\u001b[39m tvm\u001b[39m.\u001b[39;49mbuild(sch\u001b[39m.\u001b[39;49mmod, target\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mcuda\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m      3\u001b[0m A_np \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mrandom\u001b[39m.\u001b[39muniform(size\u001b[39m=\u001b[39m(\u001b[39m1024\u001b[39m,))\u001b[39m.\u001b[39mastype(\u001b[39m\"\u001b[39m\u001b[39mfloat32\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m      4\u001b[0m B_np \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mrandom\u001b[39m.\u001b[39muniform(size\u001b[39m=\u001b[39m(\u001b[39m1024\u001b[39m,))\u001b[39m.\u001b[39mastype(\u001b[39m\"\u001b[39m\u001b[39mfloat32\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda2/envs/mlc/lib/python3.8/site-packages/tvm/driver/build_module.py:281\u001b[0m, in \u001b[0;36mbuild\u001b[0;34m(inputs, args, target, target_host, runtime, name, binds)\u001b[0m\n\u001b[1;32m    277\u001b[0m     target_host \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mllvm\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m tvm\u001b[39m.\u001b[39mruntime\u001b[39m.\u001b[39menabled(\u001b[39m\"\u001b[39m\u001b[39mllvm\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mstackvm\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    279\u001b[0m annotated_mods, target_host \u001b[39m=\u001b[39m Target\u001b[39m.\u001b[39mcanon_target_map_and_host(annotated_mods, target_host)\n\u001b[0;32m--> 281\u001b[0m rt_mod_host \u001b[39m=\u001b[39m _driver_ffi\u001b[39m.\u001b[39;49mtir_to_runtime(annotated_mods, target_host)\n\u001b[1;32m    283\u001b[0m annotated_mods, target_host \u001b[39m=\u001b[39m Target\u001b[39m.\u001b[39mcanon_target_map_and_host(annotated_mods, target_host)\n\u001b[1;32m    285\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39misinstance\u001b[39m(target_host, Target):\n",
      "File \u001b[0;32mtvm/_ffi/_cython/./packed_func.pxi:331\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.PackedFuncBase.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mtvm/_ffi/_cython/./packed_func.pxi:262\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.FuncCall\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mtvm/_ffi/_cython/./packed_func.pxi:251\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.FuncCall3\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mtvm/_ffi/_cython/./base.pxi:181\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.CHECK_CALL\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mTVMError\u001b[0m: Traceback (most recent call last):\n  12: TVMFuncCall\n  11: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::Module (tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)>::AssignTypedLambda<tvm::__mk_TVM22::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#1}>(tvm::__mk_TVM22::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tvm::runtime::TVMRetValue)\n  10: tvm::TIRToRuntime(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target const&)\n  9: tvm::SplitMixedModule(tvm::IRModule, tvm::Target const&, tvm::Target const&)\n  8: tvm::ApplyPasses(tvm::IRModule, tvm::transform::Sequential)\n  7: tvm::transform::Pass::operator()(tvm::IRModule) const\n  6: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  5: tvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  4: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  3: tvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\n  2: _ZN3tvm7runtime13PackedFun\n  1: tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::tir::transform::VerifyMemory()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1}>(tvm::tir::transform::VerifyMemory()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1})::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  0: _ZN3tvm7runtime6deta\n  Did you forget to bind?\n    Variable `B` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n    Variable `A` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n    Variable `C` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.\n  File \"/workspace/tvm/src/tir/analysis/verify_memory.cc\", line 205\nRuntimeError: Memory verification failed with the following errors:\n# from tvm.script import tir as T\n\n@T.prim_func\ndef main(A: T.Buffer((1024,), \"float32\"), B: T.Buffer((1024,), \"float32\"), C: T.Buffer((1024,), \"float32\")):\n    T.func_attr({\"target\": T.target({\"arch\": \"sm_75\", \"host\": {\"keys\": [\"cpu\"], \"kind\": \"llvm\", \"tag\": \"\"}, \"keys\": [\"cuda\", \"gpu\"], \"kind\": \"cuda\", \"max_num_threads\": 1024, \"tag\": \"\", \"thread_warp_size\": 32}), \"tir.noalias\": T.bool(True)})\n    for i_0, i_1 in T.grid(8, 128):\n        cse_var_1: T.int32 = i_0 * 128 + i_1\n        C_1 = T.Buffer((1024,), data=C.data)\n        A_1 = T.Buffer((1024,), data=A.data)\n        B_1 = T.Buffer((1024,), data=B.data)\n        C_1[cse_var_1] = A_1[cse_var_1] + B_1[cse_var_1]"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
    "\n",
    "A_np = np.random.uniform(size=(1024,)).astype(\"float32\")\n",
    "B_np = np.random.uniform(size=(1024,)).astype(\"float32\")\n",
    "A_nd = tvm.nd.array(A_np, tvm.cuda(0))\n",
    "B_nd = tvm.nd.array(B_np, tvm.cuda(0))\n",
    "C_nd = tvm.nd.array(np.zeros((1024,), dtype=\"float32\"), tvm.cuda(0))\n",
    "\n",
    "rt_mod[\"main\"](A_nd, B_nd, C_nd)\n",
    "print(A_nd)\n",
    "print(B_nd)\n",
    "print(C_nd)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1.4. 示例：窗口求和\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_40826/1416523311.py:4: DeprecationWarning: T.Buffer[...] is deprecated, use T.Buffer(...) instead\n",
      "  def main(A: T.Buffer[(1027,), \"float32\"],\n",
      "/tmp/ipykernel_40826/1416523311.py:5: DeprecationWarning: T.Buffer[...] is deprecated, use T.Buffer(...) instead\n",
      "  B: T.Buffer[(1024,), \"float32\"]) -> None:\n",
      "<ast>:3: DeprecationWarning: T.Buffer[...] is deprecated, use T.Buffer(...) instead\n"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModuleWindowSum:\n",
    "    @T.prim_func\n",
    "    def main(A: T.Buffer[(1027,), \"float32\"],\n",
    "             B: T.Buffer[(1024,), \"float32\"]) -> None:\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i in T.grid(1024):\n",
    "            with T.block(\"C\"):\n",
    "                vi = T.axis.remap(\"S\", [i])\n",
    "                B[vi] = A[vi] + A[vi + 1] + A[vi + 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xling/miniconda2/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/home/xling/miniconda2/envs/mlc/bin/python -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1027</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">8</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;blockIdx.x&quot;</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">128</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;threadIdx.x&quot;</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi:vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">3</span>])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(B[vi])\n",
       "                    B[vi] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[vi] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">1</span>] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(MyModuleWindowSum)\n",
    "nthread = 128\n",
    "block_C = sch.get_block(\"C\")\n",
    "i,  = sch.get_loops(block=block_C)\n",
    "i0, i1 = sch.split(i, [None, nthread])\n",
    "sch.bind(i0, \"blockIdx.x\")\n",
    "sch.bind(i1, \"threadIdx.x\")\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "重要的是，在这种情况下，有数据复用的机会 (reuse opportunities)。 请注意，每个 GPU 线程块都包含所有线程都可以在块内访问的共享内存 (shared memory)。 我们使用cache_read添加一个中间阶段，将部分数据（下面绿色）缓存到共享内存上。 缓存完成后，线程可以从共享内存中读取。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xling/miniconda2/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/home/xling/miniconda2/envs/mlc/bin/python -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1027</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        A_shared <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">1027</span>,), scope<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;shared&quot;</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">8</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;blockIdx.x&quot;</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">128</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;threadIdx.x&quot;</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">for</span> ax0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">130</span>):\n",
       "                    <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;A_shared&quot;</span>):\n",
       "                        v0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1027</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
       "                        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v0])\n",
       "                        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(A_shared[v0])\n",
       "                        A_shared[v0] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v0]\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A_shared[vi:vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">3</span>])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(B[vi])\n",
       "                    B[vi] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_shared[vi] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A_shared[vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">1</span>] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A_shared[vi <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A_shared = sch.cache_read(block_C, read_buffer_index=0, storage_scope=\"shared\")\n",
    "sch.compute_at(A_shared, i1)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \\\n",
      "     (__CUDACC_VER_MAJOR__ > 11))\n",
      "#define TVM_ENABLE_L2_PREFETCH 1\n",
      "#else\n",
      "#define TVM_ENABLE_L2_PREFETCH 0\n",
      "#endif\n",
      "\n",
      "#ifdef _WIN32\n",
      "  using uint = unsigned int;\n",
      "  using uchar = unsigned char;\n",
      "  using ushort = unsigned short;\n",
      "  using int64_t = long long;\n",
      "  using uint64_t = unsigned long long;\n",
      "#else\n",
      "  #define uint unsigned int\n",
      "  #define uchar unsigned char\n",
      "  #define ushort unsigned short\n",
      "  #define int64_t long long\n",
      "  #define uint64_t unsigned long long\n",
      "#endif\n",
      "extern \"C\" __global__ void __launch_bounds__(128) main_kernel(float* __restrict__ A, float* __restrict__ B) {\n",
      "  __shared__ float A_shared[130];\n",
      "  for (int ax0 = 0; ax0 < 130; ++ax0) {\n",
      "    A_shared[ax0] = A[((((int)blockIdx.x) * 128) + ax0)];\n",
      "  }\n",
      "  __syncthreads();\n",
      "  B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = ((A_shared[((int)threadIdx.x)] + A_shared[(((int)threadIdx.x) + 1)]) + A_shared[(((int)threadIdx.x) + 2)]);\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
    "print(rt_mod.imported_modules[0].get_source())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1.4.1. 为其他 GPU 平台构建代码¶\n",
    "MLC 过程通常支持针对多种硬件平台，我们可以通过改变目标参数来生成 Metal 代码（这是另一种 GPU 编程模型）。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "// Function: main_kernel\n",
      "#include <metal_stdlib>\n",
      "using namespace metal;\n",
      "\n",
      "union __TVMArgUnion {\n",
      " int v_int[2];\n",
      "};\n",
      "\n",
      "kernel void main_kernel(  device float* A [[ buffer(0) ]],\n",
      "  device float* B [[ buffer(1) ]],\n",
      "  uint blockIdx [[threadgroup_position_in_grid]],\n",
      "  uint threadIdx [[thread_position_in_threadgroup]]\n",
      ") {\n",
      "  threadgroup float A_shared[130];\n",
      "  for (int ax0 = 0; ax0 < 130; ++ax0) {\n",
      "    A_shared[ax0] = A[((((int)blockIdx) * 128) + ax0)];\n",
      "  }\n",
      "  threadgroup_barrier(mem_flags::mem_threadgroup);\n",
      "  B[((((int)blockIdx) * 128) + ((int)threadIdx))] = ((A_shared[((int)threadIdx)] + A_shared[(((int)threadIdx) + 1)]) + A_shared[(((int)threadIdx) + 2)]);\n",
      "}\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[00:31:35] /workspace/tvm/src/target/opt/build_metal_off.cc:32: Warning: Metal runtime not enabled, return a source module...\n"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(sch.mod, target=\"metal\")\n",
    "print(rt_mod.imported_modules[0].get_source())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1.5. 矩阵乘法\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModuleMatmul:\n",
    "    @T.prim_func\n",
    "    def main(A: T.Buffer((1024, 1024), \"float32\"),\n",
    "             B: T.Buffer((1024, 1024), \"float32\"),\n",
    "             C: T.Buffer((1024, 1024), \"float32\")) -> None:\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i, j, k in T.grid(1024, 1024, 1024):\n",
    "            with T.block(\"C\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    C[vi, vj] = 0.0\n",
    "                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1.5.1. 本地存储分块 (Local Blocking)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xling/miniconda2/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/home/xling/miniconda2/envs/mlc/bin/python -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        C_local <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), scope<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;local&quot;</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">16</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;blockIdx.y&quot;</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">16</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;blockIdx.x&quot;</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">for</span> i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">8</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;threadIdx.y&quot;</span>):\n",
       "                    <span style=\"color: #008000; font-weight: bold\">for</span> j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>thread_binding(<span style=\"color: #008000\">8</span>, thread<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;threadIdx.x&quot;</span>):\n",
       "                        <span style=\"color: #008000; font-weight: bold\">for</span> i_2_init, j_2_init <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>):\n",
       "                            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C_init&quot;</span>):\n",
       "                                vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2_init)\n",
       "                                vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_2_init)\n",
       "                                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "                                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C_local[vi, vj])\n",
       "                                C_local[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                        <span style=\"color: #008000; font-weight: bold\">for</span> k_0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">256</span>):\n",
       "                            <span style=\"color: #008000; font-weight: bold\">for</span> k_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>unroll(<span style=\"color: #008000\">4</span>):\n",
       "                                <span style=\"color: #008000; font-weight: bold\">for</span> i_2, j_2 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>):\n",
       "                                    <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C_update&quot;</span>):\n",
       "                                        vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2)\n",
       "                                        vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_2)\n",
       "                                        vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">1024</span>, k_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> k_1)\n",
       "                                        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(C_local[vi, vj], A[vi, vk], B[vk, vj])\n",
       "                                        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C_local[vi, vj])\n",
       "                                        C_local[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C_local[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
       "                        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>):\n",
       "                            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C_local&quot;</span>):\n",
       "                                v0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
       "                                v1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax1)\n",
       "                                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(C_local[v0, v1])\n",
       "                                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v0, v1])\n",
       "                                C[v0, v1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C_local[v0, v1]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def blocking(sch,\n",
    "             tile_local_y,\n",
    "             tile_local_x,\n",
    "             tile_block_y,\n",
    "             tile_block_x,\n",
    "             tile_k):\n",
    "    block_C = sch.get_block(\"C\")\n",
    "    C_local = sch.cache_write(block_C, 0, \"local\")\n",
    "\n",
    "    i, j, k = sch.get_loops(block=block_C)\n",
    "\n",
    "    i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])\n",
    "    j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])\n",
    "    k0, k1 = sch.split(loop=k, factors=[None, tile_k])\n",
    "    sch.unroll(k1)\n",
    "    sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)\n",
    "    sch.reverse_compute_at(C_local, j1)\n",
    "\n",
    "    sch.bind(i0, \"blockIdx.y\")\n",
    "    sch.bind(j0, \"blockIdx.x\")\n",
    "\n",
    "    sch.bind(i1, \"threadIdx.y\")\n",
    "    sch.bind(j1, \"threadIdx.x\")\n",
    "    sch.decompose_reduction(block_C, k0)\n",
    "\n",
    "    return sch\n",
    "\n",
    "sch = tvm.tir.Schedule(MyModuleMatmul)\n",
    "sch = blocking(sch, 8, 8, 8, 8, 4)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GEMM-Blocking: 3359.743672 GFLOPS\n"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
    "dev = tvm.cuda(0)\n",
    "A_np = np.random.uniform(size=(1024, 1024)).astype(\"float32\")\n",
    "B_np = np.random.uniform(size=(1024, 1024)).astype(\"float32\")\n",
    "A_nd = tvm.nd.array(A_np, dev)\n",
    "B_nd = tvm.nd.array(B_np, dev)\n",
    "C_nd = tvm.nd.array(np.zeros((1024, 1024), dtype=\"float32\"), dev)\n",
    "\n",
    "num_flop = 2 * 1024 * 1024 * 1024\n",
    "evaluator = rt_mod.time_evaluator(\"main\", dev, number=10)\n",
    "\n",
    "print(\"GEMM-Blocking: %f GFLOPS\" % (num_flop / evaluator(A_nd, B_nd, C_nd).mean / 1e9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.1.7. 利用自动程序优化\n",
    "到目前为止，我们一直在手动编写变换来优化 GPU 上的 TensorIR 程序。我们可以利用自动程序优化框架来调整相同的程序。下面的代码就是这样做的，我们这里设置了一个较小的搜索次数，可能需要几分钟才能完成。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-28 00:47:57 [INFO] Logging directory: ./tune_tmp/logs\n",
      "2023-08-28 00:48:18 [INFO] LocalBuilder: max_workers = 48\n",
      "2023-08-28 00:48:19 [INFO] LocalRunner: max_workers = 1\n",
      "2023-08-28 00:48:19 [INFO] [task_scheduler.cc:159] Initializing Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>2147483648</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name          FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    2147483648         1               N/A             N/A    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                     N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "2023-08-28 00:48:19 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |       FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 2147483648 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "\n",
      "2023-08-28 00:48:19 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2023-08-28 00:48:43 [INFO] [task_scheduler.cc:193] Sending 64 sample(s) to builder\n",
      "2023-08-28 00:49:22 [INFO] [task_scheduler.cc:195] Sending 64 sample(s) to runner\n",
      "2023-08-28 00:49:28 [DEBUG] XGB iter   0: tr-p-rmse: 0.203320\ttr-a-peak@32: 1.000000\ttr-rmse: 0.699170\ttr-rmse: 0.699170\n",
      "2023-08-28 00:49:28 [DEBUG] XGB iter  25: tr-p-rmse: 0.000000\ttr-a-peak@32: 1.000000\ttr-rmse: 0.750000\ttr-rmse: 0.750000\n",
      "2023-08-28 00:49:28 [DEBUG] XGB iter  50: tr-p-rmse: 0.000000\ttr-a-peak@32: 1.000000\ttr-rmse: 0.750000\ttr-rmse: 0.750000\n",
      "2023-08-28 00:49:28 [DEBUG] XGB stopped. Best iteration: [9] tr-p-rmse:0.00000\ttr-a-peak@32:1.00000\ttr-rmse:0.75000\ttr-rmse:0.75000 \n",
      "2023-08-28 00:49:28 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>2147483648</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>64</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name          FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    2147483648         1               N/A             N/A    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                     N/A        64           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "2023-08-28 00:49:28 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |       FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 2147483648 |      1 |            N/A |          N/A |                   N/A |     64 |      \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "\n",
      "2023-08-28 00:49:28 [INFO] [task_scheduler.cc:260] Task #0 has finished. Remaining task(s): 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>2147483648</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>64</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name          FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    2147483648         1               N/A             N/A    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                     N/A        64       Y   "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-28 00:49:28 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |       FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 2147483648 |      1 |            N/A |          N/A |                   N/A |     64 |    Y \n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'NoneType' object has no attribute 'mod'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[39], line 12\u001b[0m\n\u001b[1;32m      3\u001b[0m database \u001b[39m=\u001b[39m ms\u001b[39m.\u001b[39mtune_tir(\n\u001b[1;32m      4\u001b[0m     mod\u001b[39m=\u001b[39mMyModuleMatmul,\n\u001b[1;32m      5\u001b[0m     target\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mnvidia/tesla-p100\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[39m# task_name=\"main\"\u001b[39;00m\n\u001b[1;32m     10\u001b[0m )\n\u001b[1;32m     11\u001b[0m sch \u001b[39m=\u001b[39m ms\u001b[39m.\u001b[39mtir_integration\u001b[39m.\u001b[39mcompile_tir(database, MyModuleMatmul, \u001b[39m\"\u001b[39m\u001b[39mnvidia/tesla-p100\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m---> 12\u001b[0m sch\u001b[39m.\u001b[39;49mmod\u001b[39m.\u001b[39mshow()\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'NoneType' object has no attribute 'mod'"
     ]
    }
   ],
   "source": [
    "from tvm import meta_schedule as ms\n",
    "\n",
    "database = ms.tune_tir(\n",
    "    mod=MyModuleMatmul,\n",
    "    target=\"nvidia/tesla-p100\",\n",
    "    max_trials_global=64,\n",
    "    num_trials_per_iter=64,\n",
    "    work_dir=\"./tune_tmp\",\n",
    "    # task_name=\"main\"\n",
    ")\n",
    "sch = ms.tir_integration.compile_tir(database, MyModuleMatmul, \"nvidia/tesla-p100\")\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rt_mod = tvm.build(sch.mod, target=\"nvidia/tesla-p100\")\n",
    "dev = tvm.cuda(0)\n",
    "evaluator = rt_mod.time_evaluator(\"main\", dev, number=10)\n",
    "\n",
    "print(\"MetaSchedule: %f GFLOPS\" % (num_flop / evaluator(A_nd, B_nd, C_nd).mean / 1e9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.2. 第二部分\n",
    "我们在过去的章节中讨论了为 CPU 和 GPU 环境构建机器学习编译流程。本节重点介绍我们如何为专门的硬件后端构建概念编程模型。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accel_fill_zero(C):\n",
    "    C[:] = 0\n",
    "\n",
    "def accel_tmm_add(C, A, B):\n",
    "    C[:] += A @ B.T\n",
    "\n",
    "def accel_dma_copy(reg, dram):\n",
    "    reg[:] = dram[:]\n",
    "\n",
    "def lnumpy_tmm(A: np.ndarray, B: np.ndarray, C: np.ndarray):\n",
    "    # a special accumulator memory\n",
    "    C_accumulator = np.empty((16, 16), dtype=\"float32\")\n",
    "    A_reg = np.empty((16, 16), dtype=\"float32\")\n",
    "    B_reg = np.empty((16, 16), dtype=\"float32\")\n",
    "\n",
    "    for i in range(64):\n",
    "        for j in range(64):\n",
    "            accel_fill_zero(C_accumulator[:,:])\n",
    "            for k in range(64):\n",
    "                accel_dma_copy(A_reg[:], A[i * 16 : i * 16 + 16, k * 16 : k * 16 + 16])\n",
    "                accel_dma_copy(B_reg[:], B[j * 16 : j * 16 + 16, k * 16 : k * 16 + 16])\n",
    "                accel_tmm_add(C_accumulator[:,:], A_reg, B_reg)\n",
    "            accel_dma_copy(C[i * 16 : i * 16 + 16, j * 16 : j * 16 + 16], C_accumulator[:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6.2.2.2. 带有张量化计算的 block\n",
    "我们的主要观察之一是专用加速器代码不是以标量计算为单位构建的。到目前为止，我们运行的大多数 TensorIR 代码都包含一个 block，用于计算输出张量中的单个元素。许多专门的加速器在张量区域上运行计算。 TensorIR 中的 block 结构帮助我们对此类相关计算进行分组。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MatmulBlockModule:\n",
    "    @T.prim_func\n",
    "    def main(\n",
    "        A: T.Buffer((1024, 1024), \"float32\"),\n",
    "        B: T.Buffer((1024, 1024), \"float32\"),\n",
    "        C: T.Buffer((1024, 1024), \"float32\"),\n",
    "    ) -> None:\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i0, j0, k0 in T.grid(64, 64, 64):\n",
    "            with T.block(\"tmm-16x16\"):\n",
    "                vi0, vj0, vk0 = T.axis.remap(\"SSR\", [i0, j0, k0])\n",
    "                with T.init():\n",
    "                    for i1, j1 in T.grid(16, 16):\n",
    "                        with T.block(\"tmm_init\"):\n",
    "                            vi1, vj1 = T.axis.remap(\"SS\", [i1, j1])\n",
    "                            C[vi0 * 16 + vi1, vj0 * 16 + vj1] = T.float32(0)\n",
    "\n",
    "                for i1, j1, k1 in T.grid(16, 16, 16):\n",
    "                    with T.block(\"tmm\"):\n",
    "                        vi1, vj1, vk1 = T.axis.remap(\"SSR\", [i1, j1, k1])\n",
    "                        C[vi0 *16 + vi1, vj0 * 16 + vj1] += \\\n",
    "                            A[vi0 * 16 + vi1, vk0 * 16 + vk1] * B[vj0 * 16 + vj1, vk0 * 16 + vk1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MatmulBlockModule.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "这个 block 从 A 和 B 的 16x16 区域读取，并写入 C 的 16x16 区域。 在这种情况下，block 的内容包含有关子区域计算的特定实现的更多细节。 我们将此 block 称为张量化 block，因为它们包含跨越张量子区域的计算。\n",
    "\n",
    "我们可以运行以下代码来确认 TensorIR 模块产生了正确的结果。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "\n",
    "c_nd = tvm.nd.empty((1024, 1024), dtype=\"float32\")\n",
    "\n",
    "lib = tvm.build(MatmulBlockModule, target=\"llvm\")\n",
    "lib[\"main\"](a_nd, b_nd, c_nd)\n",
    "np.testing.assert_allclose(c_nd.numpy(), c_tmm, rtol=1e-5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlc",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
